

#include "yolodef.h"
#include "parser.h"
#include "data.h"


void train_swag(char* cfgfile, char* weightfile)
{
    char* train_images = "data/voc.0712.trainval";
    char* backup_directory = "backup/";
    srand(time(0));
    char* base = basecfg(cfgfile);
    printf("%s\n", base);
    float avg_loss = -1;
    network net = parse_network_cfg(cfgfile);
    if (weightfile)
        load_weights(&net, weightfile);
    printf("Learning Rate: %g, Momentum: %g, Decay: %g\n", net.learning_rate, net.momentum, net.decay);
    int imgs = net.batch * net.subdivisions;
    int i = *net.seen / imgs;
    data train, buffer;
    layer l = net.layers[net.n - 1];
    int side = l.side;
    int classes = l.classes;
    float jitter = l.jitter;
    list* plist = get_paths(train_images);
    //int N = plist->size;
    char** paths = (char**)list_to_array(plist);
    load_args args = {0};
    args.w = net.w;
    args.h = net.h;
    args.paths = paths;
    args.n = imgs;
    args.m = plist->size;
    args.classes = classes;
    args.jitter = jitter;
    args.num_boxes = side;
    args.d = &buffer;
    args.type = REGION_DATA;
    pthread_t load_thread = load_data_in_thread(args);
    clock_t time;
    //while(i*imgs < N*120){
    while (get_current_batch(net) < net.max_batches)
    {
        i += 1;
        time = clock();
        pthread_join(load_thread, 0);
        train = buffer;
        load_thread = load_data_in_thread(args);
        printf("Loaded: %lf seconds\n", sec(clock() - time));
        time = clock();
        float loss = train_network(net, train);
        if (avg_loss < 0)
            avg_loss = loss;
        avg_loss = avg_loss * .9 + loss * .1;
        printf("%d: %f, %f avg, %f rate, %lf seconds, %d images\n", i, loss, avg_loss, get_current_rate(net), sec(clock() - time), i * imgs);
        if (i % 1000 == 0 || i == 600)
        {
            char buff[256];
            sprintf(buff, "%s/%s_%d.mo", backup_directory, base, i);
            save_weights(net, buff);
        }
        free_data(train);
    }
    char buff[256];
    sprintf(buff, "%s/%s_final.mo", backup_directory, base);
    save_weights(net, buff);
}

void run_swag(int argc, char** argv)
{
    if (argc < 4)
    {
        fprintf(stderr, "usage: %s %s [train/test/valid] [cfg] [weights (optional)]\n", argv[0], argv[1]);
        return;
    }
    char* cfg = argv[3];
    char* weights = (argc > 4) ? argv[4] : 0;
    if (0 == strcmp(argv[2], "train"))
        train_swag(cfg, weights);
}
